#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
STM32 OV7670摄像头数据接收和显示程序
支持Linux平台，通过串口接收RGB565图像数据并实时显示
"""

import serial
import cv2
import numpy as np
import time
import threading
import queue
import re

class CameraReceiver:
    def __init__(self, port='/dev/ttyUSB0', baudrate=921600):
        """
        初始化摄像头接收器
        :param port: 串口设备路径
        :param baudrate: 波特率
        """
        self.port = port
        self.baudrate = baudrate
        self.serial_conn = None
        self.frame_queue = queue.Queue(maxsize=5)
        self.running = False
        self.frame_count = 0
        
        # 图像参数
        self.width = 320
        self.height = 240
        self.pixel_count = self.width * self.height
        self.data_size = self.pixel_count * 2  # RGB565每像素2字节
        
    def connect(self):
        """连接串口"""
        try:
            self.serial_conn = serial.Serial(
                port=self.port,
                baudrate=self.baudrate,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                timeout=1.0,
                xonxoff=False,
                rtscts=False,
                dsrdtr=False
            )
            print(f"成功连接到 {self.port}, 波特率: {self.baudrate}")
            return True
        except Exception as e:
            print(f"串口连接失败: {e}")
            return False
    
    def disconnect(self):
        """断开串口连接"""
        if self.serial_conn and self.serial_conn.is_open:
            self.serial_conn.close()
            print("串口连接已关闭")
    
    def rgb565_to_rgb888(self, rgb565_data):
        """
        将RGB565数据转换为RGB888格式 - 使用大端序解析(与STM32匹配)
        :param rgb565_data: RGB565二进制数据
        :return: RGB888的numpy数组
        """
        # 检查数据长度
        if len(rgb565_data) != self.data_size:
            print(f"警告: 数据长度不匹配，期望{self.data_size}字节，实际{len(rgb565_data)}字节")
            return None
        
        # 创建BGR图像数组
        bgr888 = np.zeros((self.height, self.width, 3), dtype=np.uint8)
        
        # 逐像素转换 (使用大端序，与test.py相同的方式)
        for i in range(self.pixel_count):
            byte1 = rgb565_data[i * 2]     # 高字节
            byte2 = rgb565_data[i * 2 + 1] # 低字节
            rgb565 = (byte1 << 8) | byte2   # 大端序组合
            
            # 提取RGB分量
            r = ((rgb565 >> 11) & 0x1F) * 255 // 31  # 5位红色扩展到8位
            g = ((rgb565 >> 5) & 0x3F) * 255 // 63   # 6位绿色扩展到8位
            b = (rgb565 & 0x1F) * 255 // 31          # 5位蓝色扩展到8位
            
            # 计算像素位置 (行列坐标)
            row = i // self.width
            col = i % self.width
            
            # 存储为BGR格式 (OpenCV使用BGR)
            bgr888[row, col] = [b, g, r]
        
        return bgr888
    
    def read_line_timeout(self, timeout=2.0):
        """
        读取一行数据，带超时
        :param timeout: 超时时间（秒）
        :return: 读取到的行数据
        """
        start_time = time.time()
        line = b''
        
        while time.time() - start_time < timeout:
            if self.serial_conn.in_waiting > 0:
                char = self.serial_conn.read(1)
                if char == b'\n':
                    return line.decode('utf-8', errors='ignore').strip()
                line += char
        
        return None
    
    def receive_frame(self):
        """接收一帧图像数据 - 完整协议处理"""
        try:
            # 1. 等待帧开始标识
            while self.running:
                line = self.read_line_timeout()
                if line and line.startswith('FRAME_START:'):
                    frame_num = int(line.split(':')[1])
                    print(f"开始接收第 {frame_num} 帧")
                    break
            else:
                return None
            
            # # 2. 读取图像尺寸信息 (恢复)
            # size_line = self.read_line_timeout()
            # if not size_line or not size_line.startswith('SIZE:'):
            #     print(f"尺寸信息错误: {size_line}")
            #     # 不返回None，继续尝试读取
            # else:
            #     print(f"尺寸信息: {size_line}")
            
            # # 3. 读取格式信息 (恢复)
            # format_line = self.read_line_timeout()
            # if not format_line or not format_line.startswith('FORMAT:'):
            #     print(f"格式信息错误: {format_line}")
            #     # 不返回None，继续尝试读取
            # else:
            #     print(f"格式信息: {format_line}")
            
            # 4. 等待数据开始标记
            data_start = self.read_line_timeout()
            if not data_start or data_start != 'DATA_START':
                print(f"数据开始标记错误: {data_start}")
                return None
            
            # 5. 读取二进制图像数据 (提高效率)
            print(f"开始读取 {self.data_size} 字节的图像数据...")
            image_data = b''
            bytes_to_read = self.data_size
            
            start_time = time.time()
            chunk_size = 4096  # 增大读取块大小提高效率
            
            while len(image_data) < self.data_size and time.time() - start_time < 8.0:
                remaining = self.data_size - len(image_data)
                read_size = min(chunk_size, remaining)
                
                if self.serial_conn.in_waiting > 0:
                    chunk = self.serial_conn.read(read_size)
                    if chunk:
                        image_data += chunk
                        
                        # 每接收一定数据量显示进度
                        if len(image_data) % 30000 == 0:
                            progress = (len(image_data) / self.data_size) * 100
                            print(f"接收进度: {progress:.1f}%")
            
            if len(image_data) != self.data_size:
                print(f"图像数据大小错误: 期望 {self.data_size}, 实际 {len(image_data)}")
                return None
            
            print("图像数据接收完成")
            
            # 6. 等待数据结束标记
            while self.running:
                line = self.read_line_timeout(1.0)  # 较短超时
                if line == 'DATA_END':
                    print("找到DATA_END标记")
                    break
                elif line and 'DATA_END' in line:
                    print(f"找到包含DATA_END的行: {line}")
                    break
            
            # 7. 等待帧结束标记
            frame_end = self.read_line_timeout(1.0)
            if frame_end and frame_end.startswith('FRAME_END:'):
                end_frame_num = int(frame_end.split(':')[1])
                if end_frame_num == frame_num:
                    print(f"第 {frame_num} 帧接收完成 ✓")
                else:
                    print(f"帧号不匹配: 开始{frame_num} vs 结束{end_frame_num}")
            else:
                print(f"帧结束标记异常: {frame_end}")
            
            return image_data
            
        except Exception as e:
            print(f"接收帧数据时出错: {e}")
            return None
    
    def receiver_thread(self):
        """接收线程函数"""
        while self.running:
            frame_data = self.receive_frame()
            if frame_data:
                try:
                    # 转换图像格式
                    bgr_image = self.rgb565_to_rgb888(frame_data)
                    
                    if bgr_image is not None:
                        # 将图像放入队列
                        if not self.frame_queue.full():
                            self.frame_queue.put(bgr_image)
                            self.frame_count += 1
                        else:
                            # 队列满了，丢弃最旧的帧
                            try:
                                self.frame_queue.get_nowait()
                            except:
                                pass
                            self.frame_queue.put(bgr_image)
                            self.frame_count += 1
                        
                        print(f"✅ 成功转换并缓存第 {self.frame_count} 帧")
                    else:
                        print("❌ 图像转换失败")
                        
                except Exception as e:
                    print(f"❌ 图像转换错误: {e}")
            else:
                print("❌ 接收帧数据失败，重新尝试...")
                time.sleep(0.1)  # 短暂等待后重试
    
    def start_receive(self):
        """开始接收数据"""
        if not self.connect():
            return False
        
        self.running = True
        self.receiver_thread_handle = threading.Thread(target=self.receiver_thread)
        self.receiver_thread_handle.daemon = True
        self.receiver_thread_handle.start()
        
        print("开始接收摄像头数据...")
        return True
    
    def stop_receive(self):
        """停止接收数据"""
        self.running = False
        if hasattr(self, 'receiver_thread_handle'):
            self.receiver_thread_handle.join(timeout=2.0)
        self.disconnect()
    
    def display_frames(self):
        """显示图像帧"""
        cv2.namedWindow('OV7670 Camera', cv2.WINDOW_AUTOSIZE)
        
        fps_counter = 0
        fps_start_time = time.time()
        
        print("\n=== 控制说明 ===")
        print("'q' 或 ESC - 退出程序")
        print("'s' - 保存当前图像")
        print("================\n")
        
        while self.running:
            try:
                # 从队列获取图像
                if not self.frame_queue.empty():
                    image = self.frame_queue.get()
                    
                    # 检查图像是否有效
                    if image is None or image.size == 0:
                        print("⚠️ 获取到空图像")
                        continue
                    
                    # 计算FPS
                    fps_counter += 1
                    if fps_counter % 10 == 0:
                        fps = 10 / (time.time() - fps_start_time)
                        print(f"FPS: {fps:.2f}, 总帧数: {self.frame_count}")
                        fps_start_time = time.time()
                    
                    # 在图像上添加帧计数信息
                    cv2.putText(image, f'Frame: {self.frame_count}', 
                              (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
                    
                    # 显示图像
                    cv2.imshow('OV7670 Camera', image)
                    
                    # 检查按键
                    key = cv2.waitKey(1) & 0xFF
                    if key == ord('q') or key == 27:  # 'q'键或ESC键退出
                        break
                    elif key == ord('s'):  # 's'键保存图像
                        filename = f'captured_frame_{int(time.time())}.jpg'
                        cv2.imwrite(filename, image)
                        print(f"图像已保存为: {filename}")
                
                else:
                    time.sleep(0.01)  # 避免CPU占用过高
                    
            except KeyboardInterrupt:
                break
            except Exception as e:
                print(f"显示图像时出错: {e}")
                time.sleep(0.1)
        
        cv2.destroyAllWindows()

def main():
    """主函数"""
    print("STM32 OV7670摄像头数据接收程序")
    print("使用说明:")
    print("- 按 'q' 或 ESC 键退出")
    print("- 按 's' 键保存当前图像")
    print("-" * 50)
    
    # 创建接收器实例
    receiver = CameraReceiver(port='/dev/ttyUSB0', baudrate=921600)
    
    try:
        # 开始接收数据
        if receiver.start_receive():
            # 显示图像
            receiver.display_frames()
        else:
            print("无法启动数据接收")
    
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    
    finally:
        # 清理资源
        receiver.stop_receive()
        print("程序结束")

if __name__ == "__main__":
    main() 